#include "slice2cube.h"
#include "config-fnt.h"
#include "public/utils.h"
#include <fstream>
#include <cstring>
#include <sys/stat.h>

template <typename T>
void Slice2Cube<T>::writeCubes(Slices* slices, int jobi) {
	uint64_t blocknx=(width+xsize-1)/xsize;
	uint64_t blockny=(height+ysize-1)/ysize;
	uint64_t writern=blocknx*blockny*channels;
	for(uint64_t writeri=jobi; writeri<writern; writeri+=slices->total) {
		int channel=writeri/blocknx/blockny;
		uint64_t ystart=(writeri%(blocknx*blockny))/blocknx*ysize;
		uint64_t hh=ysize;
		if(ystart+hh>height)
			hh=height-ystart;
		uint64_t xstart=(writeri%(blocknx*blockny))%blocknx*xsize;
		uint64_t ww=xsize;
		if(xstart+ww>width)
			ww=width-xstart;
		auto& writer=writers[writeri];
		std::ofstream fs{};
		for(size_t si=0; si<slices->bufs.size(); si++) {
			auto z=slices->zidx+si;
			if(z%zsize==0) {
				std::ostringstream oss;
				oss<<outdir<<"/ch";
				oss.width(2);
				oss.fill('0');
				oss<<channel<<"/z";
				oss.width(8);
				oss.fill('0');
				oss<<z;
				oss<<"/y";
				oss.width(8);
				oss.fill('0');
				oss<<ystart<<".x";
				oss.width(8);
				oss.fill('0');
				oss<<xstart<<".nrrd";
				writer.fn=oss.str();

				fs.open(writer.fn, std::ofstream::out|std::ofstream::binary|std::ofstream::trunc);
				if(!fs)
					throwError("Failed to open");
				uint64_t dd=zsize;
				if(z+dd>depth)
					dd=depth-z;
				fs<<"NRRD0004\n";
				fs<<"# Created by " PACKAGE_NAME_LONG " " PACKAGE_VERSION "\n";
				if(sizeof(T)==1) {
					fs<<"type: uchar\n";
				} else {
					fs<<"type: ushort\n";
				}
				fs<<"encoding: gzip\n";
				fs<<"dimension: 3\n";
				fs<<"sizes: "<<ww<<" "<<hh<<" "<<dd<<"\n";
				fs<<"space: right-anterior-superior\n";
				fs<<"space directions: ("<<xres<<",0,0) (0,"<<yres<<",0) (0,0,"<<zres<<")\n";
				if(isLittleEndian()) {
					fs<<"endian: little\n";
				} else {
					fs<<"endian: big\n";
				}
				fs<<"\n";
				if(!fs)
					throwError("Failed to write header");
				writer.strm.next_out=reinterpret_cast<unsigned char*>(writer.buf);
				writer.strm.avail_out=BUFSIZ;
			}

			auto ptr=slices->bufs[si]+channel*width*height+ystart*width;
			for(uint64_t y=0; y<hh; y++) {
				writer.strm.next_in=reinterpret_cast<unsigned char*>(ptr+y*width+xstart);
				writer.strm.avail_in=ww*sizeof(T);
				do {
					if(writer.strm.avail_out==0) {
						if(!fs.is_open()) {
							fs.open(writer.fn, std::ofstream::out|std::ofstream::binary|std::ofstream::app);
							if(!fs)
								throwError("Failed to open");
						}
						fs.write(writer.buf, BUFSIZ);
						if(!fs)
							throwError("Failed to write");
						writer.strm.next_out=reinterpret_cast<unsigned char*>(writer.buf);
						writer.strm.avail_out=BUFSIZ;
					}
					auto r=deflate(&writer.strm, Z_NO_FLUSH);
					switch(r) {
						case Z_OK:
							break;
						case Z_BUF_ERROR:
							break;
						case Z_STREAM_END:
						default:
							throwError("failed to deflate");
					}
				} while(writer.strm.avail_in>0);
			}

			if((z+1)%zsize==0 || z+1==depth) {
				bool flushed=false;
				do {
					if(writer.strm.avail_out==0) {
						if(!fs.is_open()) {
							fs.open(writer.fn, std::ofstream::out|std::ofstream::binary|std::ofstream::app);
							if(!fs)
								throwError("Failed to open");
						}
						fs.write(writer.buf, BUFSIZ);
						if(!fs)
							throwError("Failed to write");
						writer.strm.next_out=reinterpret_cast<unsigned char*>(writer.buf);
						writer.strm.avail_out=BUFSIZ;
					}
					auto r=deflate(&writer.strm, Z_FINISH);
					switch(r) {
						case Z_OK:
							break;
						case Z_STREAM_END:
							flushed=true;
							break;
						case Z_BUF_ERROR:
							break;
						default:
							throwError("failed to deflate");
					}
				} while(!flushed);
				if(!fs.is_open()) {
					fs.open(writer.fn, std::ofstream::out|std::ofstream::binary|std::ofstream::app);
					if(!fs)
						throwError("Failed to open");
				}
				fs.write(writer.buf, BUFSIZ-writer.strm.avail_out);
				if(!fs)
					throwError("Failed to write");
				fs.flush();
				if(!fs)
					throwError("Failed to flush");
				fs.close();
				if(!fs)
					throwError("Failed to close");
				if(Z_OK!=deflateReset(&writer.strm)) {
					throwError("Failed to reset deflate stream.");
				}
			}
		}
		if(fs.is_open()) {
			fs.write(writer.buf, BUFSIZ-writer.strm.avail_out);
			if(!fs)
				throwError("Failed to write");
			fs.flush();
			if(!fs)
				throwError("Failed to flush");
			fs.close();
			if(!fs)
				throwError("Failed to close");
			writer.strm.next_out=reinterpret_cast<unsigned char*>(writer.buf);
			writer.strm.avail_out=BUFSIZ;
		}
	}
}
template void Slice2Cube<uint8_t>::writeCubes(Slices* slices, int jobi);
template void Slice2Cube<uint16_t>::writeCubes(Slices* slices, int jobi);

template<typename T> static bool calc_ds_mean(T* buf, T* bufds, uint64_t* bufdssum, uint64_t z, uint64_t dsw, uint64_t dsh, uint32_t dsx, uint32_t dsy, uint32_t dsz, uint64_t width, uint64_t height, uint64_t depth) {
	bool ds_first=false;
	bool ds_last=false;
	if(z%dsz==0)
		ds_first=true;
	if((z+1)%dsz==0 || z+1==depth)
		ds_last=true;
	uint64_t thickness=z%dsz+1;
	for(uint32_t y=0; y<dsh; y++) {
		uint32_t h=dsy;
		if(y*dsy+h>height)
			h=height-y*dsy;
		for(uint32_t x=0; x<dsw; x++) {
			uint32_t w=dsx;
			if(x*dsx+w>width)
				w=width-x*dsx;
			uint64_t val=0;
			if(!ds_first)
				val=bufdssum[x+y*dsw];
			auto bufImg=buf+x*dsx+y*dsy*width;
			for(uint32_t jj=0; jj<h; jj++) {
				for(uint32_t ii=0; ii<w; ii++) {
					val+=bufImg[ii+width*jj];
				}
			}
			if(ds_last) {
				auto npxl=w*h*thickness;
				bufds[x+y*dsw]=(val+npxl/2)/npxl;
			} else {
				bufdssum[x+y*dsw]=val;
			}
		}
	}
	return ds_last;
}
template<typename T> static bool calc_ds_max(T* buf, T* bufds, uint64_t* bufdssum, uint64_t z, uint64_t dsw, uint64_t dsh, uint32_t dsx, uint32_t dsy, uint32_t dsz, uint64_t width, uint64_t height, uint64_t depth) {
	bool ds_first=false;
	bool ds_last=false;
	if(z%dsz==0)
		ds_first=true;
	if((z+1)%dsz==0 || z+1==depth)
		ds_last=true;
	//uint64_t thickness=z%dsz+1;
	for(uint32_t y=0; y<dsh; y++) {
		uint32_t h=dsy;
		if(y*dsy+h>height)
			h=height-y*dsy;
		for(uint32_t x=0; x<dsw; x++) {
			uint32_t w=dsx;
			if(x*dsx+w>width)
				w=width-x*dsx;
			uint64_t val=0;
			if(!ds_first)
				val=bufdssum[x+y*dsw];
			auto bufImg=buf+x*dsx+y*dsy*width;
			for(uint32_t jj=0; jj<h; jj++) {
				for(uint32_t ii=0; ii<w; ii++) {
					if(val<bufImg[ii+width*jj])
						val=bufImg[ii+width*jj];
				}
			}
			if(ds_last) {
				bufds[x+y*dsw]=val;
			} else {
				bufdssum[x+y*dsw]=val;
			}
		}
	}
	return ds_last;
}
template<typename T> static bool calc_ds_1(T* buf, T* bufds, uint64_t width, uint64_t height) {
	for(uint32_t y=0; y<height; y++) {
		for(uint32_t x=0; x<width; x++) {
			bufds[x+y*width]=buf[x+y*width];
		}
	}
	return true;
}

template <typename T>
void Slice2Cube<T>::writeDownsample(Slices* slices, int jobi) {
	int channel_start=jobi*channels/slices->total;
	int channel_end=(jobi+1)*channels/slices->total;
	for(int i=channel_start; i<channel_end; i++) {
		auto& fsi=tempfs[i];
		for(size_t si=0; si<slices->bufs.size(); si++) {
			auto z=slices->zidx+si;
			if(z<start_idx_ds)
				continue;
			bool ds_last;
			if(dsx==1 && dsy==1 && dsz==1) {
				ds_last=calc_ds_max(slices->bufs[si]+i*width*height, &bufds[i*dsw*dsh], &bufdssum[i*dsw*dsh], z, dsw, dsh, dsx, dsy, dsz, width, height, depth);
			} else if(ds_max) {
				ds_last=calc_ds_max(slices->bufs[si]+i*width*height, &bufds[i*dsw*dsh], &bufdssum[i*dsw*dsh], z, dsw, dsh, dsx, dsy, dsz, width, height, depth);
			} else {
				ds_last=calc_ds_mean(slices->bufs[si]+i*width*height, &bufds[i*dsw*dsh], &bufdssum[i*dsw*dsh], z, dsw, dsh, dsx, dsy, dsz, width, height, depth);
			}
			if(ds_last) {
				fsi.write(reinterpret_cast<char*>(&bufds[i*dsw*dsh]), dsw*dsh*sizeof(T));
				if(!fsi)
					throwError("Failed to write image.");
				fsi.flush();
				if(!fsi)
					throwError("Failed to write image.");
			}
		}
	}
}
template void Slice2Cube<uint8_t>::writeDownsample(Slices* slices, int jobi);
template void Slice2Cube<uint16_t>::writeDownsample(Slices* slices, int jobi);

template <typename T>
void Slice2Cube<T>::deflateDownsample(Slices* slices, int jobi) {
	int channel_start=jobi*channels/slices->total;
	int channel_end=(jobi+1)*channels/slices->total;
	for(int i=channel_start; i<channel_end; i++) {
		OutputStream fs1;
		std::string fn1{tempfn[i], 0, tempfn[i].size()-4};
		fn1+="nrrd";
		if(!fs1.open(fn1))
			throwError("Failed to open file: ", fn1, ".");
		fs1<<"NRRD0004\n";
		fs1<<"# Created by " PACKAGE_NAME_LONG " " PACKAGE_VERSION "\n";
		if(sizeof(T)==1) {
			fs1<<"type: uchar\n";
		} else {
			fs1<<"type: ushort\n";
		}
		fs1<<"encoding: gzip\n";
		fs1<<"dimension: 3\n";
		fs1<<"sizes: "<<dsw<<" "<<dsh<<" "<<dsd<<"\n";
		fs1<<"space: right-anterior-superior\n";
		fs1<<"space directions: ("<<xres*dsx<<",0,0) (0,"<<yres*dsy<<",0) (0,0,"<<zres*dsz<<")\n";
		if(isLittleEndian()) {
			fs1<<"endian: little\n";
		} else {
			fs1<<"endian: big\n";
		}
		fs1<<"\n";
		if(!fs1)
			throwError("Failed to write header.");
		if(!fs1.pushGzip())
			throwError("Failed to start gzip compressing.");

		std::ifstream fs0;
		fs0.open(tempfn[i], std::ifstream::in|std::ifstream::binary);
		if(!fs0)
			throwError("Failed to open file: ");
		for(uint64_t z=0; z<dsd; z++) {
			fs0.read(reinterpret_cast<char*>(&bufds[i*dsw*dsh]), dsw*dsh*sizeof(T));
			if(!fs0)
				throwError("Failed to read file: ");
			fs1.write(reinterpret_cast<char*>(&bufds[i*dsw*dsh]), dsw*dsh*sizeof(T));
			if(!fs1)
				throwError("Failed to write file: ");
		}
		if(!fs1.pop())
			throwError("Failed to finish gzip compressing.");
		if(!fs1.close())
			throwError("Failed to close file.");
	}
}
template void Slice2Cube<uint8_t>::deflateDownsample(Slices* slices, int jobi);
template void Slice2Cube<uint16_t>::deflateDownsample(Slices* slices, int jobi);


template <typename T>
Slice2Cube<T>::NrrdBlock::NrrdBlock():
	strm{}, buf{}, fn{}
{
	strm.next_in=nullptr;
	strm.avail_in=0;
	strm.next_out=nullptr;
	strm.avail_out=0;
	strm.zalloc=nullptr;
	strm.zfree=nullptr;
	strm.opaque=nullptr;
	auto r=deflateInit2(&strm, 5, Z_DEFLATED, 15+16, 9, Z_FILTERED);
	if(r!=Z_OK)
		throwError("Failed to init gzip: ", r);
}
template Slice2Cube<uint8_t>::NrrdBlock::NrrdBlock();
template Slice2Cube<uint16_t>::NrrdBlock::NrrdBlock();

template <typename T>
Slice2Cube<T>::NrrdBlock::~NrrdBlock() {
	deflateEnd(&strm);
}
template Slice2Cube<uint8_t>::NrrdBlock::~NrrdBlock();
template Slice2Cube<uint16_t>::NrrdBlock::~NrrdBlock();
